import cv2 as cv
import numpy as np
from scipy import signal


def roberts(img, boundary='fill', fillvalue=0):
    h, w = img.shape[:2]
    h_k, w_k = 2, 2
    # 卷积核1及锚点的位置
    r1 = np.array([[1,0],[0,-1]], np.float32)
    kr1, kc1 = 0, 0
    # 计算full卷积
    img_conv_r1 = signal.convolve2d(img, r1, mode='full', boundary=boundary, fillvalue=fillvalue)
    # 根据锚点的位置截取 full 卷积，获得 same 卷积
    img_conv_r1 = img_conv_r1[h_k-kr1-1:h+h_k-kr1-1, w_k-kc1-1:w+w_k-kc1-1]
    # 卷积核2及锚点的位置
    r2 = np.array([[0, 1], [-1, 0]], np.float32)
    kr2, kc2 = 0, 1
    # 计算full卷积
    img_conv_r2 = signal.convolve2d(img, r2, mode='full', boundary=boundary, fillvalue=fillvalue)
    # 根据锚点的位置截取 full 卷积，获得 same 卷积
    img_conv_r2 = img_conv_r2[h_k - kr2 - 1:h + h_k - kr2 - 1, w_k - kc2 - 1:w + w_k - kc2 - 1]
    return img_conv_r1, img_conv_r2

if __name__ == '__main__':
    img = cv.imread("/Users/vine/Desktop/2021091120354001191G-CL0615-6.jpg", 0)
    cv.imshow('src', img)
    img_conv_r1, img_conv_r2 = roberts(img)
    # 45 方向上的边缘强度的灰度级显示
    img_conv_r1 = np.abs(img_conv_r1)
    edge_45 = img_conv_r1.astype(np.uint8)
    cv.imshow('edge_45', edge_45)
    # 135 方向上的边缘强度
    img_conv_r2 = np.abs(img_conv_r2)
    edge_135 = img_conv_r2.astype(np.uint8)
    cv.imshow('edge_135', edge_135)
    # 用平方和的开方来衡量最后输出的边缘
    edge = np.sqrt(np.power(img_conv_r1, 2.0) + np.power(img_conv_r2, 2.0))
    edge = np.round(edge)
    edge[edge>255] = 255
    edge = edge.astype(np.uint8)
    # 显示边缘
    cv.imshow('edge', edge)
    cv.waitKey(0)
    cv.destroyAllWindows()
